from cartopy import crs as ccrs
from cartopy import feature as cfeature
from matplotlib import pyplot as plt
from LiMeteo import ERA5Reader, consoleLogger
import pyvista as pv
import numpy as np


class vistaPlotter:
    def __init__(self, filePath) -> None:
        self.msg = consoleLogger('vistaPlotter')
        self.reader = ERA5Reader(filePath, msg=self.msg)
        # self.reader.printInfo()

    def setRange(self, plotRange, strTime, timeFmt="%Y%m%d%H"):
        """
        设置绘图范围
        @param plotRange:`list, tuple` 四元素数组
        @strTime:`str` 时间
        """
        self.plotRange = plotRange
        self.timeIndex = self.reader.convertTimeBack(strTime, timeFmt)

    def plotWindSpaceAnimation(self, lowRes=4, amp=20, space=30):
        levels = [400,450,500,550,600,650,700,750,775,800,825,850,875,900,925,950,975,1000]
        strTime = self.reader.time
        u = self.reader.getSubset('u', 0, 40,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3],levels, lowRes=lowRes)
        v = self.reader.getSubset('v', 0, 40,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes)
        w = self.reader.getSubset('w', 0, 40,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes) 
        lon, lat = self.reader.getSubsetLonLat(
            self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], lowRes=lowRes)
        z = np.arange(0, len(u), 1)

        wlim = 1
        w[w>wlim] = wlim
        w[w<-wlim] = -wlim
        w*= -amp

        pv.global_theme.background = 'k'
        pv.global_theme.axes.show = True
        pv.global_theme.font.color = 'white'
        grid = pv.UniformGrid()
        print(u.shape)
        grid.dimensions = np.array((u.shape[3], u.shape[2], u.shape[1]))
        gridSpace = 25
        grid.spacing = (gridSpace, gridSpace, gridSpace)
        grid.origin = (0, 0, 0)
        # pdata = np.arange(u.shape[2]* u.shape[1]* u.shape[0])
        lon *= space
        lat *= space
        z *= space

        pl = pv.Plotter()
        # add map
        plane = pv.Plane(center=(u.shape[3]*gridSpace/2, u.shape[2]*gridSpace/2, 0), i_size=u.shape[3]*gridSpace, j_size=u.shape[2]*gridSpace, direction=(0,0,-1))
        mapTex = pv.read_texture("1.png")
        pl.add_mesh(plane, texture=mapTex, lighting=False)

        timeIndex = 0
        wnd = []
        for i in np.flip(range(u.shape[1])):
            for j in np.flip(np.arange(u.shape[2])):
                for k in range(u.shape[3]):
                    wnd.append([u[timeIndex][i][j][k], v[timeIndex][i][j][k], w[timeIndex][i][j][k]])
                    # wnd.append([u[i][j][k], v[i][j][k], 0])

        # coor = np.array(coor)
        wnd = np.array(wnd)
        grid['wind'] = wnd
        grid.set_active_vectors("wind")

        grid['wind'][grid['wind'] < 8] = 0

        # grid.arrows.plot()
        pl.add_mesh(grid.arrows,  lighting=False, cmap='rainbow', name='wind')
        pl.add_text("Lilidream "+self.reader.convert_time(strTime[timeIndex]), font_size=16, name="title")
        # movie 
        pl.open_movie("1.mp4")

        tick=0
        for cameraPhi in np.arange(0, 2*np.pi, np.pi/90):
            tick+=1
            if tick%6 == 0:
                timeIndex += 1
                wnd = []
                for i in np.flip(range(u.shape[1])):
                    for j in np.flip(np.arange(u.shape[2])):
                        for k in range(u.shape[3]):
                            wnd.append([u[timeIndex][i][j][k], v[timeIndex][i][j][k], w[timeIndex][i][j][k]])
                            # wnd.append([u[i][j][k], v[i][j][k], 0])

                # coor = np.array(coor)
                wnd = np.array(wnd)
                grid['wind'] = wnd
                grid.set_active_vectors("wind")

                grid['wind'][grid['wind'] < 8] = 0

                pl.add_text("Lilidream "+self.reader.convert_time(strTime[timeIndex]), font_size=16, name="title")
                pl.add_mesh(grid.arrows,  lighting=False, cmap='rainbow', name='wind')


            # set camera
            center = np.array([u.shape[3]*gridSpace/2, u.shape[2]*gridSpace/2, u.shape[0]*gridSpace/2])
            cameraR = u.shape[3]*gridSpace*1.5
            
            # pl.show(auto_close=False)

            cameraLoc = [np.cos(cameraPhi)*cameraR+center[0], np.sin(cameraPhi)*cameraR+center[1], 25*gridSpace+np.cos(cameraPhi)*10*gridSpace]

            pl.camera.position = cameraLoc
            pl.camera.focal_point = center
            # pl.cam    
            # pl.show()
            pl.write_frame()
        pl.close()

    def plotWindSpace2(self, lowRes=4, amp=20, space=30):
        levels = [400,450,500,550,600,650,700,750,775,800,825,850,875,900,925,950,975,1000]
        u = self.reader.getSubset('u', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3],levels, lowRes=lowRes)
        v = self.reader.getSubset('v', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes)
        w = self.reader.getSubset('w', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes) 
        lon, lat = self.reader.getSubsetLonLat(
            self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], lowRes=lowRes)
        z = np.arange(0, len(u), 1)

        wlim = 1
        w[w>wlim] = wlim
        w[w<-wlim] = -wlim
        w*= -amp

        pv.global_theme.background = 'k'
        pv.global_theme.axes.show = True
        pv.global_theme.font.color = 'white'
        grid = pv.UniformGrid()
        print(u.shape)
        grid.dimensions = np.array((u.shape[2], u.shape[1], u.shape[0]))
        gridSpace = 25
        grid.spacing = (gridSpace, gridSpace, gridSpace)
        grid.origin = (0, 0, 0)
        # pdata = np.arange(u.shape[2]* u.shape[1]* u.shape[0])
        lon *= space
        lat *= space
        z *= space

        wnd = []
        for i in np.flip(range(u.shape[0])):
            for j in np.flip(np.arange(u.shape[1])):
                for k in range(u.shape[2]):
                    wnd.append([u[i][j][k], v[i][j][k], w[i][j][k]])
                    # wnd.append([u[i][j][k], v[i][j][k], 0])

        # coor = np.array(coor)
        wnd = np.array(wnd)
        grid['wind'] = wnd
        grid.set_active_vectors("wind")

        grid['wind'][grid['wind'] < 8] = 0

        # add map
        plane = pv.Plane(center=(u.shape[2]*gridSpace/2, u.shape[1]*gridSpace/2, 0), i_size=u.shape[2]*gridSpace, j_size=u.shape[1]*gridSpace, direction=(0,0,-1))
        mapTex = pv.read_texture("1.png")

        # grid.arrows.plot()
        pl = pv.Plotter()
        pl.add_mesh(grid.arrows,  lighting=False, cmap='rainbow')
        pl.add_mesh(plane, texture=mapTex, lighting=False)
        pl.add_text("Lilidream", font_size=16)


        # center = np.array([u.shape[2]*gridSpace/2, u.shape[1]*gridSpace/2, u.shape[0]*gridSpace/2])
        # cameraR = u.shape[2]*gridSpace*1.5
        
        # pl.show(auto_close=False)

        pl.show() 
        pl.close()
    
    def generateMap(self, mapRange=None, width=16):
        if mapRange is None:
            mapRange = self.plotRange
        ratio = (mapRange[3]-mapRange[2])/(mapRange[1]-mapRange[0])
        fig = plt.figure(figsize=(width-1.3, width*ratio))
        plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
        ax = fig.add_subplot(111, projection=ccrs.PlateCarree())
        ax.add_feature( cfeature.LAND, facecolor="#444444")
        ax.add_feature( cfeature.OCEAN, facecolor="#222222")
        # ax.coastlines( linewidth=1, color='k')
        ax.set_extent(mapRange)
        plt.savefig("1.png")
        # plt.show()
        

    def example(self):
        from pyvista import examples
        mesh = examples.download_carotid().threshold(145, scalars="scalars")
        print(mesh.scalars)
        # mask = mesh['scalars'] < 210
        # mesh['scalars'][mask] = 0  # null out smaller vectors

        # # Make a geometric object to use as the glyph
        # geom = pv.Arrow()  # This could be any dataset

        # # Perform the glyph
        # glyphs = mesh.glyph(orient="vectors", scale="scalars", factor=0.003, geom=geom)

        # plot using the plotting class
        pl = pv.Plotter()
        pl.add_mesh(glyphs, show_scalar_bar=False, lighting=False, cmap='coolwarm')
        pl.add_mesh(mesh)
        pl.camera_position = [
            (146.53, 91.28, 21.70),
            (125.00, 94.45, 19.81),
            (-0.086, 0.007, 0.996),
        ]  # view only part of the vector field
        cpos = pl.show(return_cpos=True)

    def plotWindSpace(self, lowRes=4, amp=20, space=30):
        """
        画静止风场用于分析
        """
        # 设置显示层数
        levels = [500,550,600,650,700,750,775,800,825,850,875,900,925,950,975,1000]
        # levels = [500, 550]

        # 格式化显示当前时间
        strTime = self.reader.time

        # 获取风场数据
        u = self.reader.getSubset('u', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3],levels, lowRes=lowRes)
        v = self.reader.getSubset('v', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes)
        w = self.reader.getSubset('w', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes) 
        # lon, lat = self.reader.getSubsetLonLat(
            # self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], lowRes=lowRes)
        hgt = self.reader.getSubset('z', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], [500], lowRes=lowRes)/98
        t = self.reader.getSubset('t', self.timeIndex, self.timeIndex+1,
                                  self.plotRange[0], self.plotRange[1], self.plotRange[2], self.plotRange[3], levels, lowRes=lowRes)-273.15
        
        # 定义各层在3d空间中的高度
        # z = np.arange(0, len(u), 1)

        # 垂直速度最大值限制
        wlim = 2
        w[w>wlim] = wlim
        w[w< -wlim] = -wlim
        w*= -amp

        # 设置绘图主题
        pv.global_theme.background = 'k'
        pv.global_theme.axes.show = True
        pv.global_theme.font.color = 'white'

        # 创建风场mesh
        grid = pv.UniformGrid()
        grid.dimensions = np.array((u.shape[2], u.shape[1], u.shape[0]))

        # 格点空间间隔
        gridSpace = 25
        grid.spacing = (gridSpace, gridSpace, gridSpace*2/lowRes)
        grid.origin = (0, 0, 0)

        # 创建高度场grid
        gridH = pv.UniformGrid()
        gridH.dimensions = np.array((hgt.shape[1], hgt.shape[0], 1))
        gridH.spacing = (gridSpace, gridSpace, gridSpace)
        gridH.origin = (0, 0, len(u)*gridSpace*2/lowRes)

        pl = pv.Plotter()
        # add map
        plane = pv.Plane(center=(u.shape[2]*gridSpace/2, u.shape[1]*gridSpace/2, 0), i_size=u.shape[2]*gridSpace, j_size=u.shape[1]*gridSpace, direction=(0,0,-1))
        mapTex = pv.read_texture("1.png")
        pl.add_mesh(plane, texture=mapTex, lighting=False)

        # 生成风场数据
        wnd = []
        wwnd = []
        temp = []
        for i in np.flip(range(u.shape[0])):
            for j in np.flip(np.arange(u.shape[1])):
                for k in range(u.shape[2]):
                    wnd.append((u[i,j,k], v[i,j,k], w[i,j,k]))
                    wwnd.append(w[i,j,k])
                    temp.append(t[i,j,k])
                    
                    # if u[i,j,k]**2+v[i,j,k]**2+w[i,j,k]**2 > 0**2:
                        # wnd.append((u[i,j,k], v[i,j,k], w[i,j,k]))
                    # else:
                        # wnd.append([0,0,0])
        wnd = np.array(wnd)
        
        grid['wind'] = wnd
        grid['w'] = wwnd
        grid['t'] = temp
        grid.set_active_vectors("wind")
        grid.set_active_scalars('t')

        temperature = grid.contour(7, scalars='t')
        sliFig = grid.slice_orthogonal()

        # 去除较小的风矢量
        # grid['wind'][grid['wind'] < 8] = 0

        hgtData = []
        for j in np.flip(np.arange(hgt.shape[0])):
            for k in range(hgt.shape[1]):
                hgtData.append([hgt[j][k]])
        hgtData = np.array(hgtData);
        gridH['hgt'] = hgtData
        hgtContour = gridH.contour(np.arange(500, 600, 4));

        pl.add_mesh(grid.arrows, scalars='w', lighting=False, cmap='rainbow',clim=[-40, 40], name='wind')
        # pl.add_mesh(grid.streamlines('wind', source_radius=400, n_points=500, max_steps=60000),lighting=False)
        pl.add_mesh(hgtContour, lighting=False, name='hgt', color='white', line_width=1)
        pl.add_mesh(temperature, opacity=0.5, lighting=False, cmap='jet')
        # pl.add_mesh(sliFig, lighting=False, cmap='jet')
        pl.show()
        pl.close()

    # 生成球面模型
    def sphereMesh(self, levels, lowRes=8, levelHeightCoef=10, R=1000):
        import vtk
        plotRange = self.plotRange
        # 以 0经度0纬度为x轴
        # 创建球面模型
        lon = np.arange(plotRange[0], plotRange[1] +0.25, 0.25*lowRes)
        lat = np.arange(plotRange[3], plotRange[2] -0.25, -0.25*lowRes)
        cosTheta = np.cos(lon*np.pi/180)
        cosPhi = np.cos(lat*np.pi/180)
        sinTheta = np.sin(lon*np.pi/180)
        sinPhi = np.sin(lat*np.pi/180)
        
        # 将气压转为高度
        levels = (np.log(np.array(levels)) + 7) * levelHeightCoef
        points = []
        cells = []
        lonLen = len(lon)
        latLen = len(lat)
        levelLen = len(levels)
        planeLen = lonLen * latLen

        for i in range(levelLen):
            print(i)
            for j in range(latLen):
                for k in range(lonLen):
                    r = R + levels[i]
                    points.append([ R * cosTheta[k] * cosPhi[j],R * cosPhi[j]*sinTheta[k], R * sinPhi[j]])

        for i in range(levelLen -1):
            print(i)
            for j in range(lonLen - 1):
                for k in range(latLen - 1):
                    p = k + j * lonLen + i * planeLen
                    k = [8, p, p+1, p+1+lonLen, p+lonLen, p+planeLen, p+1+planeLen, p+1+lonLen+planeLen, p+lonLen+planeLen]
                    cells.append(k)
        cells = np.array(cells).ravel()
        celltypes = np.empty((lonLen-1)*(latLen-1)*(levelLen-1), dtype=np.uint8)
        celltypes[:] = vtk.VTK_HEXAHEDRON
        offset = np.arange(0, len(points), 9)
        grid = pv.UnstructuredGrid(cells, celltypes, np.array(points))
        grid.plot(show_edges=True)











if __name__ == "__main__":
    # p = vistaPlotter('G:/js/weatherChart/data/20220120-20220131.nc')
    p = vistaPlotter('F:/idea/weather-chart/data/20220120-20220131.nc')
    # p.setRange([80, 140, 10, 70], "2022013008")
    p.setRange([85, 180, 10, 65], "2022013008")
    # p.generateMap()
    # p.plotWindSpace()
    p.sphereMesh(levels=[500,550,600])
    # p.example()
